
#include "glodef.h"
#include "errlog.h"
#include "matrices.h"
#include "vecs.h"
#include "grid3d.h"
#include "fe.h"
#include "fe_matr.h"

#include <stdlib.h>

#define USE_CONDENSATION 1

#if !defined(USE_CONDENSATION)
#define USE_CONDENSATION    0
#endif

/*
 *  ATTENTION: Condensation only works for p1-bubble element.
 *  condense_data is an array of FLOATs with size 9*grid->NEL.
 *  Denote lM[i][j] the local matrix of element el.  Then condense_data
 *  contains the following data:
 *      condense_data[el*9 + j] == lM[4][j], for j = 0, 1, 2, 3, 4
 *      condense_data[el*9 + 5 + i] == lM[i][4] / lM[4][4], for i = 0, 1, 2, 3
 *  The first set of numbers is needed to compute the values of the bubble
 *  degrees of freedom in post-processing.  The second set are used to modify
 *  the right hand side of the system.
 */


typedef struct _fe_matr_t
{
    unsigned el_type;
    grid3d_t *grid;
    local_matr_t *local_matr;
    get_eo_t *get_eo;
    Sparse_Matrix *matr;
    Sparse_Matrix *matr_bc;
    void *condense_data;
} fe_matr_t;

unsigned get_el_type(fe_matr_p M){ return M->el_type;}


fe_matr_t *fe_matr_alloc(unsigned type, local_matr_t *lm, get_eo_t *get_eo,
        grid3d_t *grid, int symmetric, char *mattyp)
{
    fe_matr_t *newmat;
    int ndof;
    newmat=malloc(sizeof(fe_matr_t));
    if(newmat==NULL)
        goto alloc_err;
    newmat->el_type=type;
    newmat->grid = grid;
    newmat->local_matr=lm;
    newmat->get_eo=get_eo;
    ndof=fe_ndof(grid, type);
    if(type==fe_bub && USE_CONDENSATION)
    {
        double *cd;
        newmat->matr=calloc(1, sizeof(Sparse_Matrix));
        if(!CreateSM(newmat->matr, grid->NN, grid->NN, symmetric, mattyp))
            goto alloc_err;
        newmat->matr_bc=calloc(1, sizeof(Sparse_Matrix));
        if(!CreateSM(newmat->matr_bc, grid->NN, grid->NN, FALSE, NULL))
            goto alloc_err;
        cd = malloc(9*grid->NEL*sizeof(double ));
        if(cd==NULL)
            goto alloc_err;
        newmat->condense_data = cd;
    }
    else if(type==fe_bub_vec && USE_CONDENSATION )
    {
        double *cd;
        newmat->matr=calloc(1, sizeof(Sparse_Matrix));
        if(!CreateSM(newmat->matr, 3*grid->NN, 3*grid->NN, symmetric, mattyp))
            goto alloc_err;
        newmat->matr_bc=calloc(1, sizeof(Sparse_Matrix));
        if(!CreateSM(newmat->matr_bc, 3*grid->NN, 3*grid->NN, FALSE, NULL))
            goto alloc_err;
        cd = malloc(90*grid->NEL*sizeof(double ));
        if(cd==NULL)
            goto alloc_err;
        newmat->condense_data = cd;
    }
    else
    {
        newmat->matr=calloc(1, sizeof(Sparse_Matrix));
        if(!CreateSM(newmat->matr, ndof, ndof, symmetric, mattyp))
            goto alloc_err;
        newmat->matr_bc=calloc(1, sizeof(Sparse_Matrix));
        if(!CreateSM(newmat->matr_bc, ndof, ndof, FALSE, NULL))
            goto alloc_err;
        newmat->condense_data = NULL;
    }
    return newmat;
    
    alloc_err:
        error_msg("fe_matr_alloc: memory allocation error.\n");
}



void fe_matr_free(fe_matr_t *matr)
{
    if(matr)
    {
        DestroySM(matr->matr);
        free(matr->matr);
        DestroySM(matr->matr_bc);
        free(matr->matr_bc);
        if(matr->condense_data)
            free(matr->condense_data);
        free(matr);
    }
}



static void entry_oper(fe_matr_t *m, int ni, int nj, unsigned eoi, unsigned eoj, double val)
{
    Sparse_Matrix *sm, *bm;
    sm=m->matr;
    bm=m->matr_bc;
    if((eoi==eo_ignore) || (eoj==eo_ignore))
        return;
    if((eoi==eo_bc) || (eoj==eo_bc))
    {
        if(ni==nj)
            SM_SetEntry(sm, ni, nj, 1.0);
        switch(eoi)
        {
            case eo_add:
                SM_AddToEntry(bm, ni, nj, val);
                break;
            case eo_set:
                SM_SetEntry(bm, ni, nj, val);
                break;
            default: break;
        }
        return;
    }
    switch (eoi)
    {
        case eo_add:
            switch (eoj)
            {
                case eo_add:
                    if((!sm->Symmetric)||(ni<=nj))
                        SM_AddToEntry(sm, ni, nj, val);
                    return;
                case eo_set:
                    goto err_setadd;
            }
        case eo_set:
            switch (eoj)
            {
                case eo_add:
                    goto err_setadd;
                case eo_set:
                    if((!sm->Symmetric)||(ni<=nj))
                        SM_SetEntry(sm, ni, nj, val);
                    return;
            }
    }
    err_setadd:
        error_msg("fe_matr_assemble: ambiguous set/add entry operation.\n");
}

void fe_matr_structure(const grid3d_t *grid, int nmatr, fe_matr_t **matr)
{
    fe_matr_t *m;
    int el, mi;
    for(mi=0;mi<nmatr;mi++)
    {
        SM_BeginStructure(matr[mi]->matr);
        SM_BeginStructure(matr[mi]->matr_bc);
    }
    for(el=0;el<grid->NEL;el++)
    {
        for(m=matr[mi=0];mi<nmatr;m=matr[++mi])
        {
            unsigned *eo;
            int *ln, nn, i, j;
            nn=fe_dofs(grid, el, m->el_type, &ln);
            m->get_eo(grid, m->el_type, ln, &eo);
            /* do condensation */
            if(m->el_type==fe_bub && m->condense_data!=NULL && USE_CONDENSATION)
            {
                nn--;
            }
            else if(m->el_type==fe_bub_vec && m->condense_data!=NULL && USE_CONDENSATION)
            {
                int k, j;
                int NN = grid->NN;
                int NEL = grid->NEL;
                static int foo[15];
                
                for(k=0;k<3;k++)
                {
                    for(j=0;j<4;j++)
                        foo[5*k+j] = NN*k+ln[j];
                    foo[5*k+4] = 3*NN+k*NEL+el;
                }
                ln = foo;
                eo[4] = eo[9] = eo[14] = eo_ignore;
            }
            
            for(i=0;i<nn;i++)
                for(j=0;j<nn;j++)
                    entry_oper(m, ln[i], ln[j], eo[i], eo[j], 0.0);
        }
    }
    for(mi=0;mi<nmatr;mi++)
    {
        SM_EndStructure(matr[mi]->matr);
        SM_EndStructure(matr[mi]->matr_bc);
    }
}


void fe_matr_assemble(const grid3d_t *grid, int nmatr, fe_matr_t **matr)
{
    fe_matr_t *m;
    int el, mi;
    clock_t tm = clock();
    
    PrintLog("Assembling %u matrices of dimensions ", nmatr);
    
    for(mi=0;mi<nmatr;mi++)
    {
        SM_BeginAssemble(matr[mi]->matr);
        SM_BeginAssemble(matr[mi]->matr_bc);
        PrintLog("%u ", matr[mi]->matr->rows);
    }
    for(el=0;el<grid->NEL;el++)
    {
        for(m=matr[mi=0];mi<nmatr;m=matr[++mi])
        /* if(matr[mi]) */
        {
            double **lM;
            unsigned *eo;
            int *ln, nn, i, j;
            m->local_matr(grid, el, m->el_type, &lM);
            nn=fe_dofs(grid, el, m->el_type, &ln);
            m->get_eo(grid, m->el_type, ln, &eo);
            /* do condensation */
            if(m->el_type==fe_bub && m->condense_data!=NULL && USE_CONDENSATION)
            {
                double *CD = m->condense_data;
                for(i=0;i<4;i++)
                {
                    double CC = lM[i][4] / lM[4][4];
                    if(eo[i]==eo_bc)
                        CC = 0.0;
                    else
                        for(j=0;j<4;j++)
                            lM[i][j] -= CC * lM[4][j];
                    CD[9*el + i] = lM[4][i];
                    CD[9*el + 5 + i] = CC;
                }
                CD[9*el + 4] = lM[4][4];
                nn--;
            }
            else if(m->el_type==fe_bub_vec && m->condense_data!=NULL && USE_CONDENSATION)
            {
                int k;
                double *CD = ((double *)m->condense_data) + 90*el;
                /* save coefficients of Gauss elimination of cols 4, 9, 14 */
                for(k=0;k<3;k++)
                    for(i=0;i<15;i++)
                    {
                    double Coef;
                    if(i == 4+5*k)      /* no elimination in condensation rows */
                        Coef = 0.0;
                    else if (eo[i]==eo_bc)  /* no elimination in rows of Dirchlet BC */
                        Coef = 0.0;
                    else                /* Gaussian elimination */
                    {
                        Coef = lM[i][4+5*k] / lM[4+5*k][4+5*k];
                        for(j=0;j<15;j++)
                            lM[i][j] -= Coef * lM[4+5*k][j];
                    }
                    CD[(3+k)*15+i] = Coef;
                    }
                /* save rows 4, 9, and 14 */
                for(k=0;k<3;k++)
                    for(i=0;i<15;i++)
                        CD[15*k + i] = lM[5*k + 4][i];
                
                int NN = grid->NN;
                int NEL = grid->NEL;
                int foo[15];
                
                for(k=0;k<3;k++)
                {
                    for(j=0;j<4;j++)
                        foo[5*k+j] = NN*k+ln[j];
                    foo[5*k+4] = 3*NN+k*NEL+el;
                }
                ln = foo;
                eo[4] = eo[9] = eo[14] = eo_ignore;
            }
            
            for(i=0;i<nn;i++)
                for(j=0;j<nn;j++)
                    entry_oper(m, ln[i], ln[j], eo[i], eo[j], lM[i][j]);
        }
    }
    PrintLog("    CPU time  %-6.3fsec.\n",
                (clock() - tm)/(double) CLOCKS_PER_SEC);
    tm = clock();
    PrintLog("Symbolic factorizations (if any) ");
    for(mi=0;mi<nmatr;mi++)
    {
        SM_EndAssemble(matr[mi]->matr);
        SM_EndAssemble(matr[mi]->matr_bc);
    }
    PrintLog("    CPU time  %-6.3fsec.\n",
                (clock() - tm)/(double) CLOCKS_PER_SEC);
}


void fe_matr_solve_scalar(const fe_matr_p matr, scalar X, scalar F, scalar T, double TOL)
{
    int i, dim;
    Sparse_Matrix *sm, *bm;
    if(matr==NULL)
        error_msg("NULL matrix passed to fe_matr_solve.\n");
    
    sm=matr->matr;
    bm=matr->matr_bc;
    dim=sm->rows;
    assert((sm->cols==dim) && (bm->rows==dim) && (bm->cols==dim));
    
    /* assert((X->N==dim) && (F->N==dim) && T->N==dim); */
    
    double *FF = get_F_scal(F), 
           *TF = get_F_scal(T),
           *XF = get_F_scal(X);
    
    clock_t tm = clock();
    PrintLog("Solving scalar equation of dim = %u\n", dim);
    
    /* pre-processing for condensation */
    if(matr->el_type==fe_bub_vec && matr->condense_data!=NULL && USE_CONDENSATION)
    {
        int el, *ln, k;
        double *CD = (double *)matr->condense_data;
        for(el=0;el<matr->grid->NEL;el++, CD+=90)
        {
            fe_dofs(matr->grid, el, fe_bub_vec, &ln);
            for(k=0;k<3;k++)
                for(i=0;i<15;i++)
                    FF[ln[i]] -= CD[(3+k)*15+i] * FF[ln[4+5*k]];
        }
        
        int NN = matr->grid->NN;
        int NEL =matr->grid->NEL;
        int NBUB = NN+NEL;
        copy_scal(F, T);
        for(k=0;k<3;k++)
        {
            for(i=0;i<NN;i++) FF[k*NN+i] = TF[k*NBUB+i];
            for(i=0;i<NEL;i++) FF[3*NN+k*NEL+i] = TF[k*NBUB+NN+i];
        }
    }    
    
    clean_scal(T);
    MultiplySparseByVector(bm, dim, FF, dim, TF);
    // for(i=0;i<dim;i++) FF[i] -= TF[i];
    axpy_scal(F, -1.0, T);
    
    /* NOTE: sparse_cgm checks if matrix has Symmetric flag up and calls
     * solve_cgm if it does or solve_bcgm if it doesn't.
     * So, we don't need to check for symmetry here.
     */
    PrintLog(" -> ");    SolveSparse(sm, XF, FF, TOL);
    
    /* post-processing for condensation */
    if(matr->el_type==fe_bub_vec && matr->condense_data!=NULL && USE_CONDENSATION)
    {
        int k, NN = matr->grid->NN;
        int NEL =matr->grid->NEL;
        int NBUB = NN+NEL;
        copy_scal(X, T);
        for(k=0;k<3;k++)
        {
            for(i=0;i<NN;i++) XF[k*NBUB+i] = TF[k*NN+i];
            for(i=0;i<NEL;i++) XF[k*NBUB+NN+i] = TF[3*NN+k*NEL+i];
        }
        
        int el, *ln;
        double *CD = (double *)matr->condense_data;
        for(el=0;el<matr->grid->NEL;el++, CD+=90)
        {
            fe_dofs(matr->grid, el, fe_bub_vec, &ln);
            for(k=0;k<3;k++)
            {
                XF[ln[5*k+4]] = FF[3*NN+k*NEL+el];
                for(i=0;i<15;i++)
                    XF[ln[5*k+4]] -= ((i==5*k+4)?0.0:CD[k*15+i])*XF[ln[i]];
                    XF[ln[5*k+4]] /= CD[k*15 + 5*k+4];
            }
        }
    }
    PrintLog("%s: CPU time  %-6.3fsec.\n", __func__, 
                (clock() - tm)/(double) CLOCKS_PER_SEC);
}

void fe_matr_solve_vector(const fe_matr_p matr, vector X, vector F, vector T, double TOL)
{
    int i, dim;
    Sparse_Matrix *sm, *bm;
    if(matr==NULL)
        error_msg("NULL matrix passed to fe_matr_solve.\n");
    
    sm=matr->matr;
    bm=matr->matr_bc;
    dim=sm->rows;
    
    assert((bm->rows==dim) && (bm->cols==dim) && (sm->cols==dim));
    
    double  *FX = get_X_vec(F),
            *FY = get_Y_vec(F),
            *FZ = get_Z_vec(F);
    double  *TX = get_X_vec(T),
            *TY = get_Y_vec(T),
            *TZ = get_Z_vec(T);
    double  *XX = get_X_vec(X),
            *XY = get_Y_vec(X),
            *XZ = get_Z_vec(X);
            
    
    clock_t tm = clock();
    PrintLog("Solving vector equation of dim = %u\n", dim);

    /* pre-processing for condensation */
    if(matr->el_type==fe_bub && matr->condense_data!=NULL && USE_CONDENSATION)
    {
        int el, *ln;
        double *CD = matr->condense_data;
        for(el=0;el<matr->grid->NEL;el++)
        {
            fe_dofs(matr->grid, el, fe_bub, &ln);
            for(i=0;i<4;i++)
            {
                FX[ln[i]] -= CD[9*el + 5 + i] * FX[ln[4]];
                FY[ln[i]] -= CD[9*el + 5 + i] * FY[ln[4]];
                FZ[ln[i]] -= CD[9*el + 5 + i] * FZ[ln[4]];
            }
        }
    }
    
    /* boundary conditions */
    clean_vec(T);  // this is very important
    MultiplySparseByVector(bm, dim, FX, dim, TX);
    MultiplySparseByVector(bm, dim, FY, dim, TY);
    MultiplySparseByVector(bm, dim, FZ, dim, TZ);
    // for(i=0;i<dim;i++)
    // {
    //     F->X[i] -= T->X[i];
    //     F->Y[i] -= T->Y[i];
    //     F->Z[i] -= T->Z[i];
    // }
    axpy_vec(F, -1.0, T);
    
    /* NOTE: sparse_cgm checks if matrix has Symmetric flag up and calls
     * solve_cgm if it does or solve_bcgm if it doesn't.
     * So, we don't need to check for symmetry here.
     */
    PrintLog(" -> ");    SolveSparse(sm, XX, FX, TOL);
    PrintLog(" -> ");    SolveSparse(sm, XY, FY, TOL);
    PrintLog(" -> ");    SolveSparse(sm, XZ, FZ, TOL);
    
    /* post-processing for condensation */
    if(matr->el_type==fe_bub && matr->condense_data!=NULL && USE_CONDENSATION)
    {
        int el, *ln;
        double *CD = matr->condense_data;
        for(el=0;el<matr->grid->NEL;el++)
        {
            fe_dofs(matr->grid, el, fe_bub, &ln);
            XX[ln[4]] = FX[ln[4]];
            XY[ln[4]] = FY[ln[4]];
            XZ[ln[4]] = FZ[ln[4]];
            for(i=0;i<4;i++)
            {
                XX[ln[4]] -= CD[9*el + i] * XX[ln[i]];
                XY[ln[4]] -= CD[9*el + i] * XY[ln[i]];
                XZ[ln[4]] -= CD[9*el + i] * XZ[ln[i]];
            }
            XX[ln[4]] /= CD[9*el + 4];
            XY[ln[4]] /= CD[9*el + 4];
            XZ[ln[4]] /= CD[9*el + 4];
        }
    }
    PrintLog("%s: CPU time  %-6.3fsec.\n", __func__, 
                (clock() - tm)/(double) CLOCKS_PER_SEC);
}




void eo_condense_bc_dir(const grid3d_t *grid, unsigned eltype, const int *dofs, unsigned **eos)
{
    eo_bc_dir(grid, eltype, dofs, eos);
    if(eltype == fe_bub)
        *eos[4] = eo_condense;
}

void eo_bc_dir(const grid3d_t *grid, unsigned eltype, const int *dofs, unsigned **eos)
{
    static unsigned _eo[MAX_LDOFS];
    int i, nn;
    if(eos) *eos=_eo;
    else return;
    nn=fe_dofs(grid, 0, eltype, NULL);
    for(i=0;i<nn;i++)
        _eo[i] = fe_dof_bcond(grid, eltype, dofs[i])==BC_DIR ? eo_bc : eo_add;
}

void eo_condense_add_all(const grid3d_t *grid, unsigned eltype, const int *dofs, unsigned **eos)
{
    eo_add_all(grid, eltype, dofs, eos);
    if(eltype == fe_bub)
        *eos[4] = eo_condense;
}

void eo_add_all(const grid3d_t *grid, unsigned eltype, const int *dofs, unsigned **eos)
{
    static unsigned _eo[MAX_LDOFS];
    int i, nn;
    (void) dofs;
    if(eos) *eos=_eo;
    else return;
    nn=fe_dofs(grid, 0, eltype, NULL);
    for(i=0;i<nn;i++)
        _eo[i] = eo_add;
}


/****************************************************************************/
#include <memory.h>

void project_grad_vec(const grid3d_t *grid, unsigned type, const fe_matr_p mass,
        vector u, vector gu0, vector gu1, vector gu2, 
        vector w, vector t)
{
    assert(mass != NULL);
    assert(mass->el_type == type);
    
    int uN = get_N_vec(u);
    scalar c = alloc_scal(uN);
    double *cF = get_F_scal(c);
    
    // c.F = u->X;
    memcpy(cF, get_X_vec(u), uN*sizeof(double));
    test_grad_e1_in_e2(grid, type, c, type, w);
    fe_matr_solve_vector(mass, gu0, w, t, 0.0);
    
    // c.F = u->Y;
    memcpy(cF, get_Y_vec(u), uN*sizeof(double));
    test_grad_e1_in_e2(grid, type, c, type, w);
    fe_matr_solve_vector(mass, gu1, w, t, 0.0);
    
    // c.F = u->Z;
    memcpy(cF, get_Z_vec(u), uN*sizeof(double));
    test_grad_e1_in_e2(grid, type, c, type, w);
    fe_matr_solve_vector(mass, gu2, w, t, 0.0);
    
    free_scal(c);
    return;
}

void project_grad_scal(const grid3d_t *grid, unsigned type, const fe_matr_p mass,
        scalar p, vector gp, vector w, vector t)
{
    assert(mass != NULL);
    assert(mass->el_type == type);

    test_grad_e1_in_e2(grid, type, p, type, w);
    fe_matr_solve_vector(mass, w, gp, t, 0.0);

    return;
}


